package com.hjs.homework.state;

import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.typeinfo.TypeHint;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.functions.sink.SinkFunction;

import java.util.ArrayList;
import java.util.List;

public class OperatorStateDemo implements SinkFunction<Tuple2<Long, Long>>, CheckpointedFunction {

    ListState<Tuple2<Long, Long>> listState;

    int threshold;

    private List<Tuple2<Long, Long>> bufferedElements;

    public OperatorStateDemo(int threshold) {
        this.threshold = threshold;
        this.bufferedElements = new ArrayList<>();
    }

    @Override
    public void snapshotState(FunctionSnapshotContext functionSnapshotContext) throws Exception {
        System.out.println("===snapshotState===");
        this.listState.clear();
        for (Tuple2<Long, Long> bufferedElement : bufferedElements) {
            listState.add(bufferedElement);
        }
    }

    @Override
    public void initializeState(FunctionInitializationContext functionInitializationContext) throws Exception {
        System.out.println("===initializeState===");
        ListStateDescriptor<Tuple2<Long, Long>> operatorDemo = new ListStateDescriptor<>(
                "operatorDemo",
                TypeInformation.of(new TypeHint<Tuple2<Long, Long>>() {
                })
        );

        listState = functionInitializationContext.getOperatorStateStore().getListState(operatorDemo);
        if (functionInitializationContext.isRestored()) {
            for (Tuple2<Long, Long> element : listState.get()) {
                bufferedElements.add(element);
            }
            System.out.println("===restored===");
        }
    }

    @Override
    public void invoke(Tuple2<Long, Long> value, Context context) throws Exception {
        bufferedElements.add(value);

        if (bufferedElements.size() == threshold) {
            for (Tuple2<Long, Long> bufferedElement : bufferedElements) {
                System.out.println("===out===:" + bufferedElement);
            }
            bufferedElements.clear();
        }
    }
}
